{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Common setup and functions used by the CW305 ECC demos."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CURRENT_BITFILE = 'original'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if TRACES != 'SIMULATED':\n",
    "    # Basic initialization:\n",
    "    scope.adc.offset = 0\n",
    "    scope.adc.basic_mode = \"rising_edge\"\n",
    "    scope.trigger.triggers = \"tio4\"\n",
    "    scope.io.tio1 = \"serial_rx\"\n",
    "    scope.io.tio2 = \"serial_tx\"\n",
    "    scope.io.hs2 = \"disabled\"\n",
    "\n",
    "    if PLATFORM == 'CWPRO':\n",
    "        scope.adc.stream_mode = True\n",
    "        scope.adc.samples = 1200000\n",
    "        target.pll.pll_outfreq_set(10E6, 1)\n",
    "        target._clksleeptime = 150\n",
    "        scope.gain.db = 30\n",
    "    elif PLATFORM == 'CWHUSKY':\n",
    "        scope.adc.stream_mode = True\n",
    "        scope.adc.samples = 1200000\n",
    "        target.pll.pll_outfreq_set(15E6, 1)\n",
    "        target._clksleeptime = 100\n",
    "        scope.gain.db = 20\n",
    "    elif PLATFORM == 'CWLITE':\n",
    "        scope.adc.samples = 24400\n",
    "        target.pll.pll_outfreq_set(50E6, 1)\n",
    "        target._clksleeptime = 30\n",
    "        scope.gain.db = 30\n",
    "\n",
    "\n",
    "    if TARGET_PLATFORM == 'CW312T_A35':\n",
    "        scope.clock.clkgen_freq = 7.37e6\n",
    "        scope.io.hs2 = 'clkgen'\n",
    "        scope.gain.db = 31\n",
    "        if PLATFORM == 'CWHUSKY':\n",
    "            scope.clock.clkgen_src = 'system'\n",
    "            scope.clock.adc_mul = 1\n",
    "            scope.clock.reset_dcms()\n",
    "        else:\n",
    "            scope.clock.adc_src = \"clkgen_x1\"\n",
    "        import time\n",
    "        time.sleep(0.1)\n",
    "        target._ss2_test_echo()\n",
    "\n",
    "    else:\n",
    "        if PLATFORM == 'CWHUSKY':\n",
    "            scope.clock.clkgen_freq = 15e6\n",
    "            scope.clock.clkgen_src = 'extclk'\n",
    "            scope.clock.adc_mul = 1\n",
    "        else:\n",
    "            scope.clock.adc_src = \"extclk_x1\"\n",
    "\n",
    "    if PLATFORM == 'CWHUSKY':\n",
    "        scope.adc.offset = 3\n",
    "    else:\n",
    "        scope.adc.offset = 0\n",
    "\n",
    "    if 'CW305' in TARGET_PLATFORM:\n",
    "        target.vccint_set(1.0)\n",
    "        # we only need PLL1:\n",
    "        target.pll.pll_enable_set(True)\n",
    "        target.pll.pll_outenable_set(False, 0)\n",
    "        target.pll.pll_outenable_set(True, 1)\n",
    "        target.pll.pll_outenable_set(False, 2)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "cycles = np.load('data/ecc_cycles.npy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def set_adc(samples):\n",
    "    if PLATFORM == 'CWPRO':\n",
    "        scope.adc.stream_mode = True\n",
    "        scope.adc.samples = samples\n",
    "        scope.adc.offset = 0\n",
    "    elif PLATFORM == 'CWHUSKY':\n",
    "        scope.adc.stream_mode = True\n",
    "        scope.adc.samples = samples\n",
    "        scope.adc.offset = 3\n",
    "        scope.adc.segments = 1\n",
    "    elif PLATFORM == 'CWLITE':\n",
    "        scope.adc.samples = 24400\n",
    "        scope.adc.offset = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def random_k(bits=256, tries=100):\n",
    "    import random\n",
    "    if TRACES == 'SIMULATED':\n",
    "        return None\n",
    "    for i in range(tries):\n",
    "        k = random.getrandbits(bits)\n",
    "        if k < target.curve.order and k > 0:\n",
    "            return k\n",
    "    raise ValueError(\"Failed to generate a valid random k after %d tries!\" % self.tries)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from chipwhisperer.common.traces import Trace\n",
    "from tqdm.notebook import trange\n",
    "import numpy as np\n",
    "import time\n",
    "import math \n",
    "\n",
    "SEGMENTS = 257 # +1 so we can grab the trailing POIs\n",
    "SEGMENT_CYCLES = 4204\n",
    "\n",
    "def get_traces(N=50, k=0, step='part1_1', randomize_k=False, full=False, samples_per_segment=256, as_int=True):\n",
    "    samples = 1130000\n",
    "    traces = []\n",
    "            \n",
    "    if TRACES == 'SIMULATED':\n",
    "        # eh maybe not optimal but it works\n",
    "        try:\n",
    "            from cwtraces import sca204_lab_data\n",
    "            load_traces = sca204_lab_data[\"get_traces\"]\n",
    "        except:\n",
    "            print(\"Simulated traces require cwtraces (pip install cwtraces)\")\n",
    "            raise\n",
    "        raws = load_traces('%s.npz' % step)\n",
    "        for t in raws['arr_0']:\n",
    "            traces.append(Trace(t[0], t[1], t[2], None))\n",
    "        raws.close()\n",
    "        print('Pre-recorded traces loaded ✅.')\n",
    "\n",
    "    else:\n",
    "        attempt4 = get_bitfile_version() == 'attempt4'\n",
    "        if PLATFORM == 'CWHUSKY':\n",
    "            scope.adc.bits_per_sample = 8 # for smaller recorded traces; doesn't appear to impact attack success rates\n",
    "        else:\n",
    "            if not full:\n",
    "                print('Warning: only \"full\" capture mode is supported with CW-lite/Pro.')\n",
    "            full = True # force full capture on non-Husky because not supported\n",
    "        if PLATFORM == 'CWLITE':\n",
    "            set_adc(samples)\n",
    "        else:\n",
    "            if not full:\n",
    "                scope.adc.segments = SEGMENTS\n",
    "                scope.adc.segment_cycles = SEGMENT_CYCLES\n",
    "                scope.adc.segment_cycle_counter_en = True\n",
    "                scope.adc.samples = samples_per_segment\n",
    "                scope.adc.stream_mode = True\n",
    "                scope.adc.offset = 3\n",
    "            else:\n",
    "                set_adc(samples)\n",
    "                if PLATFORM == 'CWHUSKY':\n",
    "                    scope.adc.segments = 1\n",
    "                    scope.adc.segment_cycles = 0\n",
    "                    scope.adc.segment_cycle_counter_en = False\n",
    "\n",
    "        for i in trange(N, desc='Capturing traces'):\n",
    "            P = target.new_point() # every trace uses a different point\n",
    "            if randomize_k:\n",
    "                k = random_k()\n",
    "            assert k != 0\n",
    "            if attempt4:\n",
    "                kb = 0x10000000000000000000000000000000000000000000000000000000000000000 - k\n",
    "                target.fpga_write(target.REG_KB, list(int.to_bytes(kb, length=32, byteorder='little')))\n",
    "\n",
    "            if PLATFORM == 'CWPRO' or PLATFORM == 'CWHUSKY':\n",
    "                ret = target.capture_trace(scope, Px=P.x, Py=P.y, k=k, check=True, as_int=as_int)\n",
    "                if not ret:\n",
    "                    print(\"Failed capture\")\n",
    "                    continue\n",
    "                traces.append(ret)\n",
    "\n",
    "            elif PLATFORM == 'CWLITE':\n",
    "                scope.adc.offset = 0\n",
    "                wave = np.array([])\n",
    "                \n",
    "                segments = math.ceil(target.pmul_cycles / scope.adc.samples)\n",
    "                for j in range(segments):\n",
    "                    ret = target.capture_trace(scope, Px=P.x, Py=P.y, k=k, check=True, as_int=as_int)\n",
    "                    if not ret:\n",
    "                        print(\"Failed capture\")\n",
    "                        continue\n",
    "                    wave = np.append(wave, ret.wave)\n",
    "                    scope.adc.offset += scope.adc.samples\n",
    "                traces.append(Trace(wave[1:], ret.textin, ret.textout, None))\n",
    "\n",
    "        if TRACES == 'COLLECT':\n",
    "            np.savez_compressed('data/%s.npz' % step, np.asarray(traces, dtype=object))\n",
    "            \n",
    "    return traces\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_sums(traces, poi):\n",
    "    sums = []\n",
    "    # in case samples were recorded as ints, translate result to make it as though they were floats\n",
    "    if 'int' in str(type(traces[0].wave[0])):\n",
    "        shift = True\n",
    "        if PLATFORM != 'CWHUSKY':\n",
    "            center = 2**9\n",
    "            div = 2**10\n",
    "        # infer whether trace was collected with 8 or 12 bits per sample:\n",
    "        elif max(abs(traces[0].wave)) > 255:\n",
    "            center = 2**11\n",
    "            div = 2**12\n",
    "        else:\n",
    "            center = 2**7\n",
    "            div = 2**8\n",
    "    else:\n",
    "        shift = False\n",
    "\n",
    "    if len(traces[0].wave) >= 1130000:\n",
    "        # full captures\n",
    "        for c in cycles:\n",
    "            sum = 0\n",
    "            for trace in traces:\n",
    "                for i in poi:\n",
    "                    power = trace.wave[c+abs(i)]\n",
    "                    if shift:\n",
    "                        power = (power-center)/div\n",
    "                    if i < 0:\n",
    "                        sum -= power\n",
    "                    else:\n",
    "                        sum += power\n",
    "            sums.append(sum/len(traces))\n",
    "    else:\n",
    "        # segmented captures (used for pre-captured traces, to save space)\n",
    "        segment_size = len(traces[0].wave) // SEGMENTS\n",
    "        for c in range(256):\n",
    "            sum = 0\n",
    "            for trace in traces:\n",
    "                for i in poi:\n",
    "                    # complicated mapping to deal with the segmented traces\n",
    "                    if abs(i) > segment_size:\n",
    "                        absi = segment_size - (SEGMENT_CYCLES - abs(i))\n",
    "                    else:\n",
    "                        absi = abs(i)\n",
    "                    if i < 0:\n",
    "                        i = -absi\n",
    "                    else:\n",
    "                        i = absi\n",
    "                    index = c*segment_size+abs(i) + cycles[0]\n",
    "                    power = trace.wave[index]\n",
    "                    if shift:\n",
    "                        power = (power-center)/div\n",
    "                    if i < 0:\n",
    "                        sum -= power\n",
    "                    else:\n",
    "                        sum += power\n",
    "            sums.append(sum/len(traces))\n",
    "\n",
    "    return sums"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_corrected_sums(traces, poi):\n",
    "    sums = []\n",
    "    # in case samples were recorded as ints, translate result to make it as though they were floats\n",
    "    if 'int' in str(type(traces[0].wave[0])):\n",
    "        shift = True\n",
    "        if PLATFORM != 'CWHUSKY':\n",
    "            center = 2**9\n",
    "            div = 2**10\n",
    "        # infer whether trace was collected with 8 or 12 bits per sample:\n",
    "        elif max(abs(traces[0].wave)) > 255:\n",
    "            center = 2**11\n",
    "            div = 2**12\n",
    "        else:\n",
    "            center = 2**7\n",
    "            div = 2**8\n",
    "    else:\n",
    "        shift = False\n",
    "        \n",
    "    if len(traces[0].wave) >= 1130000:\n",
    "        # full captures\n",
    "        for c in range(len(cycles)-1):\n",
    "            sum = 0\n",
    "            for trace in traces:\n",
    "                for p in poi:\n",
    "                    # shortcut: use the ~halfway point to determine whether the leakage influences the current bit or not\n",
    "                    if abs(p) > 2000:\n",
    "                        power = trace.wave[cycles[c]+abs(p)]\n",
    "                    else:\n",
    "                        power = trace.wave[cycles[c+1]+abs(p)]\n",
    "                    if shift:\n",
    "                        power = (power-center)/div\n",
    "                    if p < 0:\n",
    "                        sum -= power\n",
    "                    else:\n",
    "                        sum += power\n",
    "            sums.append(sum/len(traces))\n",
    "    else:\n",
    "        # segmented captures (used for pre-captured traces, to save space)\n",
    "        segment_size = len(traces[0].wave) // SEGMENTS\n",
    "        for c in range(len(cycles)-1):\n",
    "            sum = 0\n",
    "            for trace in traces:\n",
    "                for p in poi:\n",
    "                    # complicated mapping to deal with the segmented traces; also we (mis-)use segment_size to determine whether the leakage influences the current bit or not\n",
    "                    if abs(p) > segment_size:\n",
    "                        absp = segment_size - (SEGMENT_CYCLES - abs(p))\n",
    "                        d = c\n",
    "                    else:\n",
    "                        absp = abs(p)\n",
    "                        d = c + 1\n",
    "                    if p < 0:\n",
    "                        p = -absp\n",
    "                    else:\n",
    "                        p = absp\n",
    "                    index = d*segment_size+abs(p) + cycles[0]\n",
    "                    power = trace.wave[index]\n",
    "                    if shift:\n",
    "                        power = (power-center)/div\n",
    "                    if p < 0:\n",
    "                        sum -= power\n",
    "                    else:\n",
    "                        sum += power\n",
    "            sums.append(sum/len(traces))\n",
    "\n",
    "    return sums"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_corrs(traces):\n",
    "    corrsxonly = []\n",
    "    corrsyonly = []\n",
    "    corrszonly = []\n",
    "    corrsall = []\n",
    "\n",
    "    segment_size = len(traces[0].wave) // SEGMENTS\n",
    "\n",
    "    for i in range (0, len(cycles)-1):\n",
    "        corrx = 0\n",
    "        corry = 0\n",
    "        corrz = 0\n",
    "\n",
    "        if len(traces[0].wave) >= 1130000:\n",
    "            start1 = cycles[i] + rupdate_offset\n",
    "            stop1  = cycles[i] + rupdate_offset + rupdate_cycles\n",
    "\n",
    "            start2x = cycles[i+1] + rxread_offset\n",
    "            start2y = cycles[i+1] + ryread_offset\n",
    "            start2z = cycles[i+1] + rzread_offset\n",
    "\n",
    "            stop2x  = cycles[i+1] + rxread_offset + rupdate_cycles\n",
    "            stop2y  = cycles[i+1] + ryread_offset + rupdate_cycles\n",
    "            stop2z  = cycles[i+1] + rzread_offset + rupdate_cycles\n",
    "\n",
    "        else:\n",
    "            start1 = cycles[0] + (i+1)*segment_size - (SEGMENT_CYCLES - rupdate_offset)\n",
    "            stop1  = cycles[0] + (i+1)*segment_size - (SEGMENT_CYCLES - rupdate_offset) + rupdate_cycles\n",
    "\n",
    "            start2x = cycles[0] + (i+1)*segment_size + rxread_offset\n",
    "            start2y = cycles[0] + (i+1)*segment_size + ryread_offset\n",
    "            start2z = cycles[0] + (i+1)*segment_size + rzread_offset\n",
    "\n",
    "            stop2x  = cycles[0] + (i+1)*segment_size + rxread_offset + rupdate_cycles\n",
    "            stop2y  = cycles[0] + (i+1)*segment_size + ryread_offset + rupdate_cycles\n",
    "            stop2z  = cycles[0] + (i+1)*segment_size + rzread_offset + rupdate_cycles\n",
    "\n",
    "\n",
    "        for trace in traces:\n",
    "            corrx += np.corrcoef(trace.wave[start1:stop1], trace.wave[start2x:stop2x])[0][1]\n",
    "            corry += np.corrcoef(trace.wave[start1:stop1], trace.wave[start2y:stop2y])[0][1]\n",
    "            corrz += np.corrcoef(trace.wave[start1:stop1], trace.wave[start2z:stop2z])[0][1]\n",
    "        \n",
    "        #corrsall.append((corrx+corry+corrz)/len(traces))\n",
    "        # consider only Y component for attack; uncomment above to study effect of other X/Z components:\n",
    "        corrsall.append(corry/len(traces))\n",
    "        #corrsall.append((corry+corrx)/len(traces))\n",
    "    return corrsall\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def poi_guess(metric, thresholds):\n",
    "    poi_init_threshold, poi_reg_threshold = thresholds\n",
    "    guess = ''\n",
    "    if get_bitfile_version() == 'attempt4':\n",
    "        start = 1\n",
    "        initial = False\n",
    "    else:\n",
    "        start = 0\n",
    "        initial = True\n",
    "    for kbit in range(start,255):\n",
    "        if initial:\n",
    "            if metric[kbit] < poi_init_threshold:\n",
    "                guess += '0'\n",
    "            else:\n",
    "                guess += '1'\n",
    "                initial = False\n",
    "        else:\n",
    "            if metric[kbit] < poi_reg_threshold:\n",
    "                guess += '0'\n",
    "            else:\n",
    "                guess += '1'\n",
    "            \n",
    "    return guess"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def corr_guess(metric):\n",
    "    initial = True\n",
    "    guess = ''\n",
    "    for kbit in range(0,255):\n",
    "        if initial:\n",
    "            if metric[kbit] > corr_init_threshold:\n",
    "                guess += '0'\n",
    "            else:\n",
    "                guess += '1'\n",
    "                initial = False\n",
    "        else:\n",
    "            if metric[kbit] > corr_reg_threshold:\n",
    "                guess += '0'\n",
    "            else:\n",
    "                guess += '1'\n",
    "\n",
    "    return guess"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_guess(guess, k, verbose=False):\n",
    "    guesses = []\n",
    "    if get_bitfile_version() == 'attempt4':\n",
    "        top = 254\n",
    "        for a in (['0', '1']):\n",
    "            for b in (['0', '1']):\n",
    "                guesses.append(int(a + guess + b, 2))\n",
    "    else:\n",
    "        top = 255\n",
    "        guesses = [int(guess + '0', 2), int(guess + '1', 2)]\n",
    "\n",
    "    if k in guesses:\n",
    "        return ('Guessed right!', 0, 0)\n",
    "    else:\n",
    "        wrong_bits = []\n",
    "        for kbit in range(top):\n",
    "            if int(guess[kbit]) != ((k >> (top-kbit)) & 1):\n",
    "                wrong_bits.append(top-kbit)\n",
    "        if verbose:\n",
    "            print('Attack failed.')\n",
    "            print('Guesses: %s' % hex(guesses[0]))\n",
    "            for guess in guesses[1:]:\n",
    "                print('         %s' % hex(guess))\n",
    "            print('Correct: %s' % hex(k))\n",
    "            print('Wrong bits: %s' % wrong_bits)\n",
    "        return ('Failed: %3d wrong bits' % len(wrong_bits), len(wrong_bits), wrong_bits)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def poi_guess_threshold(metric, distance_threshold, thresholds):\n",
    "    poi_init_threshold, poi_reg_threshold = thresholds\n",
    "    guess = ''\n",
    "    guessed_bits = []\n",
    "    distances = []\n",
    "    \n",
    "    if get_bitfile_version() == 'attempt4':\n",
    "        start = 1\n",
    "        initial = False\n",
    "    else:\n",
    "        start = 0\n",
    "        initial = True\n",
    "    \n",
    "    if distance_threshold <= 0:\n",
    "        raise ValueEror(\"Threshold must be greater than 0\")\n",
    "        \n",
    "    #1. Calculate distances from decision thresholds:\n",
    "    for kbit in range(0,255):\n",
    "        if initial:\n",
    "            distances.append(abs(metric[kbit]- poi_init_threshold))\n",
    "        else:\n",
    "            distances.append(abs(metric[kbit]- poi_reg_threshold))\n",
    "\n",
    "    #2. Calculate the mininum distance from decision threshold for which we'll enter a guess:\n",
    "    avg = np.average(distances)\n",
    "    top = max(distances)\n",
    "    base = top-avg\n",
    "    distance_threshold = distance_threshold * base\n",
    "\n",
    "    #3. \n",
    "    if get_bitfile_version() == 'attempt4':\n",
    "        initial = False\n",
    "    else:\n",
    "        initial = True\n",
    "    for kbit in range(start,255):\n",
    "        if initial:\n",
    "            if abs(metric[kbit] - poi_init_threshold) > distance_threshold:\n",
    "                guessed_bits.append(kbit)\n",
    "            else:\n",
    "                pass\n",
    "            if metric[kbit] > poi_init_threshold:\n",
    "                guess += '0'\n",
    "            else:\n",
    "                guess += '1'\n",
    "                initial = False\n",
    "        else:\n",
    "            if abs(metric[kbit] - poi_reg_threshold) > distance_threshold:\n",
    "                guessed_bits.append(kbit)\n",
    "            else:\n",
    "                pass\n",
    "            if metric[kbit] < poi_reg_threshold:\n",
    "                guess += '0'\n",
    "            else:\n",
    "                guess += '1'\n",
    "    \n",
    "    return guess, guessed_bits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_trace_segments(N=50, poi=[-6, 7, 4202, -4203], randomize_k=False, k=0, husky_timed_segments=True, step='partXXX', as_int=True):\n",
    "    trace_segments = []\n",
    "    if TRACES == 'SIMULATED':\n",
    "        # eh maybe not optimal but it works\n",
    "        try:\n",
    "            from cwtraces import sca204_lab_data\n",
    "            load_traces = sca204_lab_data[\"get_traces\"]\n",
    "        except:\n",
    "            print(\"Simulated traces require cwtraces (pip install cwtraces)\")\n",
    "            raise\n",
    "        raws = load_traces('%s.npz' % step)\n",
    "        for t in raws['arr_0']:\n",
    "            trace_segments.append(Trace(t[0], t[1], t[2], None))\n",
    "        raws.close()\n",
    "        print('Pre-recorded traces loaded.')\n",
    "        return trace_segments\n",
    "\n",
    "    attempt4 = get_bitfile_version() == 'attempt4'\n",
    "    if PLATFORM == 'CWPRO' or (PLATFORM == 'CWHUSKY' and not husky_timed_segments): # note this approach can be used for Husky as well, but the segmented capture is faster!:\n",
    "        if PLATFORM == 'CWHUSKY':\n",
    "            scope.adc.segments = 1\n",
    "            scope.adc.segment_cycles = 0\n",
    "            scope.adc.offset = 3\n",
    "        else:\n",
    "            scope.adc.offset = 0\n",
    "        scope.adc.stream_mode = True\n",
    "        scope.adc.samples = 1120000\n",
    "        \n",
    "        for i in trange(N, desc='Capturing traces'):\n",
    "            P = target.new_point() # every trace uses a different point\n",
    "            \n",
    "            if randomize_k:\n",
    "                k = random_k()\n",
    "            assert k != 0\n",
    "            if attempt4:\n",
    "                kb = 0x10000000000000000000000000000000000000000000000000000000000000000 - k\n",
    "                target.fpga_write(target.REG_KB, list(int.to_bytes(kb, length=32, byteorder='little')))\n",
    "\n",
    "            ret = target.capture_trace(scope, Px=P.x, Py=P.y, k=k, as_int=as_int)\n",
    "            if not ret:\n",
    "                print(\"Failed capture\")\n",
    "                continue\n",
    "            trace_segment = []\n",
    "            for c in cycles:\n",
    "                for p in poi:\n",
    "                    trace_segment.append(ret.wave[c+abs(p)])\n",
    "            trace_segments.append(Trace(trace_segment, ret.textin, ret.textout, None))\n",
    "            \n",
    "    elif PLATFORM == 'CWHUSKY':\n",
    "        scope.adc.stream_mode = False\n",
    "        scope.adc.segments = 256\n",
    "        scope.adc.segment_cycles = 4204\n",
    "        scope.adc.segment_cycle_counter_en = True\n",
    "        scope.adc.samples = 11\n",
    "        scope.adc.offset = int(cycles[0] + 4201 + 3)\n",
    "        if poi == [-6, 7, 4202, -4203]:\n",
    "            indices = [1, 2, 9, 10]\n",
    "        elif poi == [-6, 7, 4201, -4202]:\n",
    "            indices = [0, 1, 9, 10]\n",
    "        elif poi == [-6, 7]:\n",
    "            indices = [9, 10]\n",
    "        else:\n",
    "            raise ValueError(\"Sorry, Husky timed segments only work for a specific set of markers; either set husky_timed_segments=False, or write your own segmented capture function\" % poi)\n",
    "\n",
    "        for i in trange(N, desc='Capturing traces'):\n",
    "            P = target.new_point() # every trace uses a different point\n",
    "            \n",
    "            if randomize_k:\n",
    "                k = random_k()\n",
    "            assert k != 0\n",
    "            if attempt4:\n",
    "                kb = 0x10000000000000000000000000000000000000000000000000000000000000000 - k\n",
    "                target.fpga_write(target.REG_KB, list(int.to_bytes(kb, length=32, byteorder='little')))\n",
    "            \n",
    "            ret = target.capture_trace(scope, Px=P.x, Py=P.y, k=k, as_int=as_int)\n",
    "            if not ret:\n",
    "                print(\"Failed capture\")\n",
    "                continue\n",
    "            trace_segment = [0, 0] # first two samples are missed but that's inconsequential since they provide no useful side channel leakage\n",
    "            for j,c in enumerate(cycles):\n",
    "                base = scope.adc.samples*j\n",
    "                for i,p in enumerate(poi):\n",
    "                    trace_segment.append(ret.wave[base+indices[i]])\n",
    "            trace_segments.append(Trace(trace_segment, ret.textin, ret.textout, None))\n",
    "\n",
    "    elif PLATFORM == 'CWLITE':\n",
    "        raise ValueError('Not implemented for CW-lite')\n",
    "\n",
    "    if TRACES == 'COLLECT':\n",
    "        np.savez_compressed('data/%s.npz' % step, np.asarray(trace_segments, dtype=object))\n",
    "\n",
    "    return trace_segments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gget_segment_sums(trace_segments, poi):\n",
    "    # Note: crucial that poi be identical to that used for get_trace_segments! (including order of elements)\n",
    "    sums = []\n",
    "    npois = len(poi)\n",
    "    for c in range(len(cycles)-1):\n",
    "        sum = 0\n",
    "        for segment in trace_segments:\n",
    "            for i,p in enumerate(poi):\n",
    "                # shortcut: use the ~halfway point to determine whether the leakage influences the current bit or not\n",
    "                if abs(p) > 2000:\n",
    "                    base = c*npois\n",
    "                else:\n",
    "                    base = (c+1)*npois\n",
    "                if p > 0:\n",
    "                    sum += segment.wave[base+i]\n",
    "                else:\n",
    "                    sum -= segment.wave[base+i]\n",
    "        sums.append(sum/len(trace_segments))\n",
    "    return sums"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_segment_sums(trace_segments, poi):\n",
    "    # Note: crucial that poi be identical to that used for get_trace_segments! (including order of elements)\n",
    "    sums = []\n",
    "    # in case samples were recorded as ints, translate result to make it as though they were floats\n",
    "    if 'int' in str(type(traces[0].wave[0])):\n",
    "        shift = True\n",
    "        if PLATFORM != 'CWHUSKY':\n",
    "            center = 2**9\n",
    "            div = 2**10\n",
    "        # infer whether trace was collected with 8 or 12 bits per sample:\n",
    "        elif max(abs(traces[0].wave)) > 255:\n",
    "            center = 2**11\n",
    "            div = 2**12\n",
    "        else:\n",
    "            center = 2**7\n",
    "            div = 2**8\n",
    "    else:\n",
    "        shift = False\n",
    "        \n",
    "    npois = len(poi)\n",
    "    for c in range(len(cycles)-1):\n",
    "        sum = 0\n",
    "        for segment in trace_segments:\n",
    "            for i,p in enumerate(poi):\n",
    "                # shortcut: use the ~halfway point to determine whether the leakage influences the current bit or not\n",
    "                if abs(p) > 2000:\n",
    "                    base = c*npois\n",
    "                else:\n",
    "                    base = (c+1)*npois\n",
    "                power = segment.wave[base+i]\n",
    "                if shift:\n",
    "                    power = (power-center)/div\n",
    "                if p > 0:\n",
    "                    sum += power\n",
    "                else:\n",
    "                    sum -= power\n",
    "        sums.append(sum/len(trace_segments))\n",
    "    return sums"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def change_bitfile(VERSION):\n",
    "    if TRACES == 'SIMULATED':\n",
    "        global CURRENT_BITFILE\n",
    "        CURRENT_BITFILE = VERSION\n",
    "    else:\n",
    "        global target\n",
    "        if VERSION == 'original':\n",
    "            rev = 0\n",
    "        elif VERSION == 'attempt1':\n",
    "            rev = 1\n",
    "        elif VERSION == 'attempt2':\n",
    "            rev = 2\n",
    "        elif VERSION == 'attempt3':\n",
    "            rev = 3\n",
    "        elif VERSION == 'attempt4':\n",
    "            rev = 4\n",
    "        else:\n",
    "            raise ValueError(\"Unsupported version %s\" % VERSION)\n",
    "        if target._fpga_id in ['cw312t_a35', '35t'] and rev == 3:\n",
    "            raise ValueError(\"attempt3 is not supported on this platform (the FPGA is not large enough for it)\")\n",
    "        target.dis()\n",
    "        target = cw.target(scope, cw.targets.CW305_ECC, force=True, fpga_id=target._fpga_id, platform=target.platform, version=rev)\n",
    "        assert get_bitfile_version() == VERSION\n",
    "\n",
    "        if PLATFORM == 'CWHUSKY':\n",
    "            # on Husky, reloading the FPGA will cause Husky's external clock frequency monitor to flag an error:\n",
    "            import time\n",
    "            time.sleep(0.5)\n",
    "            scope.errors.clear()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_bitfile_version():\n",
    "    if TRACES == 'SIMULATED':\n",
    "        return CURRENT_BITFILE\n",
    "    else:\n",
    "        rev = target.fpga_read(target.REG_CRYPT_REV, 1)[0]\n",
    "        if rev == 0:\n",
    "            return \"original\"\n",
    "        elif rev == 1:\n",
    "            return \"attempt1\"\n",
    "        elif rev == 2:\n",
    "            return \"attempt2\"\n",
    "        elif rev == 3:\n",
    "            return \"attempt3\"\n",
    "        elif rev == 4:\n",
    "            return \"attempt4\"\n",
    "        else:\n",
    "            raise ValueError(\"Warning: unrecognized version % d.\" % rev)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def consecutives(trace_segments, poi, distance_threshold, thresholds):\n",
    "\n",
    "    wrong_bits = []\n",
    "    solid_guessed_bits = []\n",
    "    total_wrong_bits = 0\n",
    "    total_solid_guessed_bits = 0\n",
    "    total_right_solid_guesses = 0\n",
    "    total_wrong_solid_guesses = 0\n",
    "    correct_solid_guesses = []\n",
    "    all_wrong_bits = []\n",
    "\n",
    "    print('Computing averages...')\n",
    "    for trace_segment in trace_segments:\n",
    "        sums = get_segment_sums([trace_segment], poi)\n",
    "\n",
    "        guess, tguessed_bits = poi_guess_threshold(sums, distance_threshold, thresholds)\n",
    "        (status, num_wrong_bits, twrong_bits) = check_guess(guess, trace_segment.textin['k'])\n",
    "\n",
    "        total_wrong_bits += num_wrong_bits\n",
    "        all_wrong_bits.append(num_wrong_bits)\n",
    "        total_solid_guessed_bits += len(tguessed_bits)\n",
    "\n",
    "        wrong_solid_guesses = len(set(twrong_bits) & set(tguessed_bits))\n",
    "        right_solid_guesses = len(tguessed_bits) - wrong_solid_guesses\n",
    "\n",
    "        total_wrong_solid_guesses += wrong_solid_guesses\n",
    "        total_right_solid_guesses += right_solid_guesses\n",
    "\n",
    "        wrong_bits.append(twrong_bits)\n",
    "        solid_guessed_bits.append(tguessed_bits)\n",
    "\n",
    "        correct_solid_guesses.append(list(set(tguessed_bits) - set(twrong_bits)))\n",
    "\n",
    "    print('All results are per-trace averages:')\n",
    "    print('Average number of wrong bits (all guesses):     %5.1f' % (total_wrong_bits/len(trace_segments)))\n",
    "    print('Average number of solid guessed bits:           %5.1f' % (total_solid_guessed_bits/len(trace_segments)))\n",
    "    print('Average number of correct solid guessed bits:   %5.1f' % (total_right_solid_guesses/len(trace_segments)))\n",
    "    print('Average number of incorrect solid guessed bits: %5.1f' % (total_wrong_solid_guesses/len(trace_segments)))\n",
    "\n",
    "    print('Computing number of good traces...')\n",
    "    # stats when taking only what we think are good guesses\n",
    "    min_c_len = 3 # we only care about at least this many correct consecutive guesses\n",
    "    total_good_consecutives = 0\n",
    "    total_bad_consecutives = 0\n",
    "    all_run_counts = np.zeros(255, np.int16)\n",
    "    good_traces = 0\n",
    "    bad_good_traces = 0\n",
    "    good_trace_ids = []\n",
    "    for t in range(len(trace_segments)):\n",
    "        run_counts = np.zeros(255, np.int16)\n",
    "        good_trace = False\n",
    "        bad_good_trace = False\n",
    "\n",
    "        # now we look for consecutive guesses, among the list of good *and* bad guesses - then we'll flag whether any bad guesses snuck in there\n",
    "        guesses = np.sort(solid_guessed_bits[t])\n",
    "        consecutives = np.split(guesses, np.where(np.diff(guesses) != 1)[0]+1)\n",
    "        good_consecutives = 0\n",
    "        bad_consecutives = 0\n",
    "        for i,c in enumerate(consecutives):\n",
    "            if len(c) >= min_c_len:\n",
    "                if any(x in consecutives[i] for x in wrong_bits[t]):\n",
    "                    bad_consecutives += 1\n",
    "                    bad_good_trace = True\n",
    "                else:\n",
    "                    good_consecutives += 1\n",
    "                    run_counts[len(c)] += 1\n",
    "                    if len(c) >= 5:\n",
    "                        good_trace = True\n",
    "        total_good_consecutives += good_consecutives\n",
    "        total_bad_consecutives += bad_consecutives\n",
    "        all_run_counts += run_counts\n",
    "        if run_counts[3] >= 3 or run_counts[4] >= 2:\n",
    "            good_trace = True\n",
    "        if good_trace:\n",
    "            good_traces += 1\n",
    "            good_trace_ids.append(t)\n",
    "            if bad_good_trace:\n",
    "                bad_good_traces += 1\n",
    "\n",
    "    print(\"Total good consecutives: %3d (%5.2f per traces)\" % (total_good_consecutives, float(total_good_consecutives/len(trace_segments))))\n",
    "    print(\"Total bad consecutives: %3d (%5.2f per traces)\" % (total_bad_consecutives, float(total_bad_consecutives/len(trace_segments))))\n",
    "    print('Number of good traces: %d' % good_traces)\n",
    "    print('Number of BAD good traces: %d' % bad_good_traces)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def find0to1trans(data):\n",
    "    pattern = [0,1]\n",
    "    return [i for i in range(0,len(data)) if list(data[i:i+len(pattern)])==pattern]\n",
    "\n",
    "def check_adc_clock_phase():\n",
    "    #scope.LA.enabled = True\n",
    "    #scope.LA.clk_source = 'target'\n",
    "    #scope.LA.oversampling_factor = 20\n",
    "    #scope.LA.capture_group = 'CW 20-pin'\n",
    "    #scope.LA.capture_depth = 50\n",
    "\n",
    "    scope.LA.arm()\n",
    "    scope.LA.trigger_now()\n",
    "\n",
    "    raw = scope.LA.read_capture_data()\n",
    "    adcclock    = scope.LA.extract(raw, 8)\n",
    "    hs1clock    = scope.LA.extract(raw, 4)\n",
    "\n",
    "    hs1_edge = find0to1trans(hs1clock)[0]\n",
    "    adc_hs1_delta = find0to1trans(adcclock[hs1_edge:])[0]\n",
    "    assert adc_hs1_delta == 13, 'Got unexpected delta: %d' % adc_hs1_delta"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
